#include "tppmc/TF2J.h"

#include <TRandom.h>
#include <TRandom3.h>
#include <TF2.h>
#include <TMath.h>

#include <unistd.h>
#include <iostream>
#include <fstream>
#include <stdlib.h>
#include <time.h>
#include <math.h>
#include <time.h>

using namespace std;

// part types if FAVORED/UNFAVORED * 20 bins in depolarization angle
#define N_DEPOL_BINS 20
#define N_PART_TYPES 2*N_DEPOL_BINS 
#define QS_NBINS 230
enum PartonType{GLUON,UQUARK,UBAR,DQUARK,DBAR,SQUARK,SBAR,CQUARK,CBAR,BQUARK,BBAR,UNKNOWN};

TF2J::TF2J(const char *name, Double_t (*fcn)(Double_t *, Double_t *), Double_t xmin, Double_t xmax, Double_t ymin, Double_t ymax, Int_t npar)
: TF2(name, fcn, xmin, xmax, ymin, ymax, npar)
{
   // allocate memory for our cache
   fIntegral_cache = new double**[N_PART_TYPES];
   for (int i = 0; i < N_PART_TYPES; ++i) {
      fIntegral_cache[i] = NULL;
      fIntegral_cache[i] = new double*[QS_NBINS];
   } // for
   
   for(int i=0; i<N_PART_TYPES; i++){
      for(int j=0; j<QS_NBINS; j++){
         fIntegral_cache[i][j] = NULL;
      } // for
   } // for
   fIntegral = NULL; 
}

TF2J::TF2J(const char *name, ROOT::Math::ParamFunctor fcn, Double_t xmin,
           Double_t xmax, Double_t ymin, Double_t ymax, Int_t npar)
: TF2(name, fcn, xmin, xmax, ymin, ymax, npar)
{
  // allocate memory for our cache
  fIntegral_cache = new double**[N_PART_TYPES];
  for (int i = 0; i < N_PART_TYPES; ++i) {
    fIntegral_cache[i] = NULL;
    fIntegral_cache[i] = new double*[QS_NBINS];
  } // for
  
  for(int i=0; i<N_PART_TYPES; i++){
    for(int j=0; j<QS_NBINS; j++){
      fIntegral_cache[i][j] = NULL;
    } // for
  } // for
  fIntegral = NULL; 
}

TF2J::~TF2J()
{
  for (int i = 0; i < N_PART_TYPES; ++i) {
    for (int j = 0; j < QS_NBINS; ++j){
      if(fIntegral_cache[i][j]){ delete [] fIntegral_cache[i][j]; }
      }
    if(fIntegral_cache[i]){ delete [] fIntegral_cache[i]; }
  }
  if(fIntegral_cache){ delete [] fIntegral_cache; }
}


void TF2J::GetRandom2(Double_t &xrandom, Double_t &yrandom)
{
  qs = pow(GetParameter(1),2); //Q^2
  if(qs>100000){
    cout << "Q^2 > 100k, setting Q^2 = 100k (TF2J)" << endl;
    qs=100000;
  }
   
  int ic = static_cast<int>(GetParameter(5));
  int qtyp = static_cast<int>(GetParameter(2));
  int fqtyp = qtyp; 
  
  if(ic==1){
    // positive pion
    if(qtyp==UQUARK) fqtyp = UQUARK; 
    if(qtyp==DQUARK) fqtyp = DQUARK;
    if(qtyp==UBAR) fqtyp = DQUARK;
    if(qtyp==DBAR) fqtyp = UQUARK; 
    if((qtyp==SBAR)||(qtyp==SQUARK)) fqtyp = DQUARK; // always unfavored
  }
  else if(ic==-1){
    // negative pion
    if(qtyp==DQUARK) fqtyp = UQUARK;
    if(qtyp==UQUARK) fqtyp = DQUARK; 
    if(qtyp==UBAR) fqtyp = UQUARK;
    if(qtyp==DBAR) fqtyp = DQUARK; 
    if((qtyp==SBAR)||(qtyp==SQUARK)) fqtyp = DQUARK; // always unfavored
  }
  else if(ic==0){
    // neutral pion
    if(qtyp==UQUARK) fqtyp = UQUARK; 
    if(qtyp==DQUARK) fqtyp = UQUARK;
    if(qtyp==UBAR) fqtyp = UQUARK;
    if(qtyp==DBAR) fqtyp = UQUARK; 

    if((qtyp==SBAR)||(qtyp==SQUARK)) fqtyp = DQUARK; // always unfavored
  }

  // Determine favored ir unfavored:
  if(fqtyp==UQUARK)
    type_j = 0;
  else if(fqtyp==DQUARK)
    type_j = 1;
 
  if(type_j!=0 && type_j!=1){ 
    cout << "type_j != 0,1, so TF2J.cc is going to segfault" << endl;
  }

  float depol_factor = GetParameter(6); 
  if( (depol_factor<0.0) || (depol_factor>1.0) ){
    cout << "woah woah woah! GetParameter(6)<0 or >1 in TF2J.cc, this is going to get messy..." << endl;
  }
  int depol_bin = int(depol_factor*(N_DEPOL_BINS-1));

  // Now, combine the FF type and depol factor to get the type_j bin

  if(type_j==0) {
    type_j += depol_bin;  
  }
  else{
    type_j = N_DEPOL_BINS + depol_bin; 
  }

  if(GetParameter(4)!=1){ 
    cout << "woah woah woah! GetParameter(4) != 1 in TF2J.cc, you're not using pi+ alone anymore! change things!" << endl;
    cout << "also, TF2J.cc going to segfault...... now!" << endl;
  }

  
  ////////////////////////// WHICH BIN DOES THIS Q^2 VALUE BELONG TO //////////////////////////
  double macro_bin_edges[] = {1, 1.25, 1.5, 2.5, 4, 6.4, 10, 15, 25, 40, 64, 100, 180, 580, 1000, 1800, 3200, 5800, 10000, 18000, 32000, 58000, 100000};
  for(int i4=0; i4<QS_NBINS; i4++){
    mod_fact = i4%10;
    lower = floor(static_cast<double>(i4)/10.0);
    micro_bin_edge = macro_bin_edges[lower+1]-macro_bin_edges[lower];
    micro_bin_edge /= 10.0;
    binsize = micro_bin_edge;
    micro_bin_edge *= mod_fact;
    micro_bin_edge += macro_bin_edges[lower];

    //cout << micro_bin_edge << endl;
    if(qs < micro_bin_edge){ 
      qs_j = i4-1;
      break;
    }
    if(qs == micro_bin_edge){
      qs_j = i4;
      break;
    }
  }
  
  //cout << "qs_j: " << qs_j << " qs: " << qs << endl;
  if(qs_j>QS_NBINS){ cout << "____________________QS_J > QS_NBINS, GOING TO SEGFAULT (TF2J.cc)__________________" << endl; }
  if(qs_j<0){ cout << "_____________________QS_J < 0, GOING TO SEGFAULT (TF2J.cc)________________" << endl; }
  
   //////////////////////////// IS THIS Q^ and Type combination already cached? if not, build it /////////////////////
   Int_t i,j,cell;
   Double_t dx   = (fXmax-fXmin)/fNpx;
   Double_t dy   = (fYmax-fYmin)/fNpy;
   Int_t ncells = fNpx*fNpy;
   
   //cout << "fIntegral_cache  [" << type_j << "][" << qs_j << "]: " << fIntegral_cache[type_j][qs_j] << endl;
   
   if (fIntegral_cache[type_j][qs_j]==NULL) {
      //cout << fIntegral_cache[type_j][qs_j] <<  endl;
      fIntegral_cache[type_j][qs_j] = new double[fNpx*fNpy+1];
      fIntegral_cache[type_j][qs_j][0] = 0.0;  // init the first cell to zero! 
      //cout << "alloced properly?[" << type_j << "][" << qs_j << "]: " << fIntegral_cache[type_j][qs_j] << endl;
      //cout << fIntegral_cache[type_j][qs_j] << "  " << fIntegral_cache[type_j][qs_j][0] << endl;
      Double_t integ;
      Int_t intNegative = 0;
      cell = 0;
      for (j=0;j<fNpy;j++) {
         for (i=0;i<fNpx;i++) {
            integ = Integral(fXmin+i*dx,fXmin+i*dx+dx,fYmin+j*dy,fYmin+j*dy+dy);
            if (integ < 0) {intNegative++; integ = -integ;}
            fIntegral_cache[type_j][qs_j][cell+1] = fIntegral_cache[type_j][qs_j][cell] + integ;
            cell++;
         }
      }
      if (intNegative > 0) {
         Warning("GetRandom2","function:%s has %d negative values: abs assumed",GetName(),intNegative);
      }
      if (fIntegral_cache[type_j][qs_j][ncells] == 0) {
         Error("GetRandom2","Integral of function is zero");
         return;
      }
      for (i=1;i<=ncells;i++) {  // normalize integral to 1
         fIntegral_cache[type_j][qs_j][i] /= fIntegral_cache[type_j][qs_j][ncells];
      }

   }
   //cout << endl;
   //////////////////////////// Set temp array (fIntegral) to be used by BinarySearch //////////////////
   fIntegral = fIntegral_cache[type_j][qs_j];
      
   /////////////////////// BINARY SEARCH AND INTERPOLATION ////////////////////////////
   Double_t r,ddx,ddy,dxint;
   r     = gRandom->Rndm();
   cell  = TMath::BinarySearch(ncells,fIntegral,r);
   dxint = fIntegral[cell+1] - fIntegral[cell];
   if (dxint > 0) ddx = dx*(r - fIntegral[cell])/dxint;
   else           ddx = 0;
   ddy = dy*gRandom->Rndm();
   j   = cell/fNpx;
   i   = cell%fNpx;
   xrandom = fXmin +dx*i +ddx;
   yrandom = fYmin +dy*j +ddy;
   
}
